import torch
from torch.nn import Module
from seasoncnn import CnnSeason
from gru import Gru


class Model(Module):
    def __init__(self, features, seq, target_id):
        super(Model, self).__init__()
        self.season_model = CnnSeason(3, 2, features, 64, seq, target_id)
        self.trend_model = Gru(features, hidden_size=64, output_size=1, num_layers=3)
        self.target_id = target_id


    def forward(self, base:torch.Tensor, trend:torch.Tensor, season:torch.Tensor):
        base = base[:,-1,self.target_id]
        season = self.season_model(season)
        trend = self.trend_model(trend)
        res = base + season + trend
        # print(res.shape)
        return res
    

    def predict(self, base:torch.Tensor, trend:torch.Tensor, season:torch.Tensor):
        self.eval()
        return self.forward(base, trend, season)